import numpy as np
import matplotlib.pyplot as plt
import numpy.random as npr
import matplotlib
matplotlib.use(backend="TkAgg")

p=0.2
trials=100_000

'''
Draw samples from the geometric distribution.

Bernoulli trials are experiments with one of two outcomes:
success or failure (an example of such an experiment is flipping
a coin).  The geometric distribution models the number of trials
that must be run in order to achieve success.  It is therefore
supported on the positive integers, ``k = 1, 2, ...``.

The probability mass function of the geometric distribution is

.. math:: f(k) = (1 - p)^{k - 1} p

where `p` is the probability of success of an individual trial.
'''
samples_geom = npr.geometric(p=p, size=trials)

theoretical_mean_geom = 1 / p

empirical_mean_geom  = sum(samples_geom) / trials

# 画图
plt.figure(figsize=(6,4))
plt.hist(samples_geom, bins=range(1, 30), density=True, alpha=0.6, color='orchid', edgecolor='black')
plt.axvline(theoretical_mean_geom, color='red', linestyle='--', label=f'Theoretical E[X]={theoretical_mean_geom:.2f}')
plt.axvline(empirical_mean_geom, color='green', linestyle=':', label=f'Empirical mean={empirical_mean_geom:.2f}')
plt.title(f'Geometric Distribution (p={p})')
plt.xlabel("X")
plt.ylabel("Frequency")
plt.legend()
plt.show()

